import os.path

import cv2
import numpy as np

imgs_dir = "/Users/john/PycharmProjects/yyzhao/drlsheep/imgs"


def match_template(image):
    # 加载模版
    template_image_paths = [os.path.join(imgs_dir, _) for _ in os.listdir(imgs_dir) if _.endswith('.png')]

    mask = np.zeros(image.shape[:2])

    threshold = 0.8

    # 可点击区域的掩码值
    mask_index = 1

    action_imgs = {}
    action_list = []

    for template_image_path in template_image_paths:
        print(f"Processing template {template_image_path}")
        template_image = cv2.imread(template_image_path)
        # 获取模板的高度和宽度
        template_image = template_image[5:-5, 5:-5]
        h, w, _ = template_image.shape

        # 使用模板匹配方法
        result = cv2.matchTemplate(image, template_image, cv2.TM_CCOEFF_NORMED)

        # 找到所有匹配的区域
        loc = np.where(result >= threshold)

        for pt in zip(*loc[::-1]):
            top = pt[0]
            bottom = pt[0] + h
            left = pt[1]
            right = pt[1] + w

            # 填充掩模
            mask[left:right, top:bottom] = mask_index
            action_imgs[mask_index] = template_image_path
            mask_index += 1

            action_list.append((top+h//2, left+w//2))

    return mask, action_list, action_imgs


def find_first_layer(image):
    # 转换到灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 计算灰度图的平均亮度
    mean_brightness = np.mean(gray)

    # 应用阈值操作
    # 第一层
    _, binary = cv2.threshold(
        gray,
        165,
        200,
        cv2.THRESH_BINARY
    )

    template_image_paths = [
        os.path.join(imgs_dir, _) for _ in os.listdir(imgs_dir) if _.endswith('.png')
    ]

    threshold = 0.8
    result = {}

    for template_image_path in template_image_paths:
        print(f"Processing template {template_image_path}")
        template_image = cv2.imread(template_image_path)
        # 获取模板的高度和宽度
        template_image = template_image[5:-5, 5:-5]
        h, w, _ = template_image.shape

        # 使用模板匹配方法
        result = cv2.matchTemplate(large_image, template_image, cv2.TM_CCOEFF_NORMED)

        # 找到所有匹配的区域
        loc = np.where(result >= threshold)

        for pt in zip(*loc[::-1]):
            top = pt[0]
            bottom = pt[0] + h
            left = pt[1]
            right = pt[1] + w

            template_name = template_image_path.split('/')[-1].split('.')[0]
            result[template_name] = {
                template_name: {
                    'center': (top+h//2, left+w//2),
                    'top': top,
                    'bottom': bottom,
                    'left': left,
                    'right': right
                }
            }


def find_first_layer_v2(image):
    # 加载模版
    template_image_paths = [os.path.join(imgs_dir, _) for _ in os.listdir(imgs_dir) if _.endswith('.png')]

    center_list = []
    threshold = 0.8
    error_distance = 10

    results = {}
    temp_index = 0

    for template_image_path in template_image_paths:
        print(f"Processing template {template_image_path}")

        template_name = template_image_path.split('/')[-1].split('.')[0]
        template_image = cv2.imread(template_image_path)

        template_image = cv2.resize(template_image, (80, 80))
        # 获取模板的高度和宽度
        h, w, _ = template_image.shape

        # template_image = template_image[5:-5, 5:-5]

        # 使用模板匹配方法
        result_img = cv2.matchTemplate(image, template_image, cv2.TM_CCOEFF_NORMED)

        # 找到所有匹配的区域
        loc = np.where(result_img >= threshold)

        for pt in zip(*loc[::-1]):
            top = pt[0]
            bottom = pt[0] + h
            left = pt[1]
            right = pt[1] + w

            center = (top+h//2, left+w//2)

            flag = True
            for center_tmp in center_list:
                if abs(center[0] - center_tmp[0]) < error_distance and abs(
                        center[1] - center_tmp[1]) < error_distance:
                    flag = False
                    break

            if flag:
                center_list.append(center)

                results[f"{template_name}_{temp_index}"] = {
                    'center': center,
                    'top': top,
                    'bottom': bottom,
                    'left': left,
                    'right': right
                }

                temp_index += 1

    return results


def find_second_layer(w, h, first_layer_result):
    # 误差距离
    error_distance = 5

    results = {}

    center_list = []

    for k, v in first_layer_result.items():
        # print(k, v)
        center = v['center']
        top = v['top']
        bottom = v['bottom']
        left = v['left']
        right = v['right']

        # 裁剪子图
        # 左上角
        center_tmp = (top, left)

        flag = True

        for center in center_list:
            if abs(center[0] - center_tmp[0]) < error_distance and abs(center[1] - center_tmp[1]) < error_distance:
                flag = False
                break

        if flag:
            top_now = top - 50
            left_now = left - 50
            bottom_now = top + 50
            right_now = left + 50

            if top_now > 0 and bottom_now < h and left_now > 0 and right_now < w:
                center_list.append(center_tmp)

                results[f"{k}_左上角"] = {
                    "top": top_now,
                    "bottom": bottom_now,
                    "left": left_now,
                    "right": right_now,
                }
        continue
        # 上方
        center_tmp = (top, (left+right)//2)

        flag = True

        for center in center_list:
            if abs(center[0] - center_tmp[0]) < error_distance and abs(center[1] - center_tmp[1]) < error_distance:
                flag = False
                break

        if flag:
            top_now = top - 50
            left_now = left
            bottom_now = top + 50
            right_now = right

            if top_now > 0 and bottom_now < h and left_now > 0 and right_now < w:
                center_list.append(center_tmp)

                results[f"{k}_左上角"] = {
                    "top": top_now,
                    "bottom": bottom_now,
                    "left": left_now,
                    "right": right_now,
                }

        # 右上角
        center_tmp = (top, right)

        flag = True

        for center in center_list:
            if abs(center[0] - center_tmp[0]) < error_distance and abs(center[1] - center_tmp[1]) < error_distance:
                flag = False
                break

        if flag:
            top_now = top - 50
            left_now = right - 50
            bottom_now = top + 50
            right_now = right + 50

            if top_now > 0 and bottom_now < h and left_now > 0 and right_now < w:
                center_list.append(center_tmp)

                results[f"{k}_左上角"] = {
                    "top": top_now,
                    "bottom": bottom_now,
                    "left": left_now,
                    "right": right_now,
                }

        # 右方
        center_tmp = ((top+bottom)//2, right)
        flag = True

        for center in center_list:
            if abs(center[0] - center_tmp[0]) < error_distance and abs(center[1] - center_tmp[1]) < error_distance:
                flag = False
                break

        if flag:
            top_now = top
            left_now = right - 50
            bottom_now = bottom
            right_now = right + 50

            if top_now > 0 and bottom_now < h and left_now > 0 and right_now < w:
                center_list.append(center_tmp)

                results[f"{k}_左上角"] = {
                    "top": top_now,
                    "bottom": bottom_now,
                    "left": left_now,
                    "right": right_now,
                }


        # 右下角
        center_tmp = (bottom, right)
        flag = True

        for center in center_list:
            if abs(center[0] - center_tmp[0]) < error_distance and abs(center[1] - center_tmp[1]) < error_distance:
                flag = False
                break
        if flag:
            top_now = bottom - 50
            left_now = right - 50
            bottom_now = bottom + 50
            right_now = right + 50

            if top_now > 0 and bottom_now < h and left_now > 0 and right_now < w:
                center_list.append(center_tmp)

                results[f"{k}_左上角"] = {
                    "top": top_now,
                    "bottom": bottom_now,
                    "left": left_now,
                    "right": right_now,
                }

        # 下方
        center_tmp = (bottom, (left+right)//2)
        flag = True
        for center in center_list:
            if abs(center[0] - center_tmp[0]) < error_distance and abs(center[1] - center_tmp[1]) < error_distance:
                flag = False
                break
        if flag:
            top_now = bottom - 50
            left_now = left
            bottom_now = bottom + 50
            right_now = right

            if top_now > 0 and bottom_now < h and left_now > 0 and right_now < w:
                center_list.append(center_tmp)

                results[f"{k}_左上角"] = {
                    "top": top_now,
                    "bottom": bottom_now,
                    "left": left_now,
                    "right": right_now,
                }

        # 左下角
        center_tmp = (bottom, left)
        flag = True
        for center in center_list:
            if abs(center[0] - center_tmp[0]) < error_distance and abs(center[1] - center_tmp[1]) < error_distance:
                flag = False
                break
        if flag:
            top_now = bottom - 50
            left_now = left - 50
            bottom_now = bottom + 50
            right_now = left + 50

            if top_now > 0 and bottom_now < h and left_now > 0 and right_now < w:
                center_list.append(center_tmp)

                results[f"{k}_左上角"] = {
                    "top": top_now,
                    "bottom": bottom_now,
                    "left": left_now,
                    "right": right_now,
                }

        # 左方
        center_tmp = ((top+bottom)//2, left)
        flag = True
        for center in center_list:
            if abs(center[0] - center_tmp[0]) < error_distance and abs(center[1] - center_tmp[1]) < error_distance:
                flag = False
                break
        if flag:
            top_now = top
            left_now = left - 50
            bottom_now = bottom
            right_now = left + 50

            if top_now > 0 and bottom_now < h and left_now > 0 and right_now < w:
                center_list.append(center_tmp)

                results[f"{k}_左上角"] = {
                    "top": top_now,
                    "bottom": bottom_now,
                    "left": left_now,
                    "right": right_now,
                }

    return results


def save_second_layer_imgs(image, first_layer_result, second_layer_result):
    pass


if __name__ == '__main__':
    template_w, template_h = 80, 80

    # 加载大图
    large_image = cv2.imread('demo_new.png')

    w, h, _ = large_image.shape

    first_layer_result = find_first_layer_v2(large_image)

    print(f"第一层检测出的元素有：{len(first_layer_result)}.")

    for k, v in first_layer_result.items():
        # print(k, v)
        cv2.rectangle(large_image, (v['top'], v['left']), (v['bottom'], v['right']), (0, 0, 255), 2)

    cv2.imwrite('first_layer.png', large_image)

    second_layer_result = find_second_layer(w, h, first_layer_result)

    print(f"第二层检测出的元素有：{len(second_layer_result)}.")
    for k, v in second_layer_result.items():
        # print(k, v)
        cv2.rectangle(large_image, (v['top'], v['left']), (v['bottom'], v['right']), (255, 0, 0), 2)

    cv2.imwrite('second_layer.png', large_image)
    # 显示结果
    # cv2.imshow('Matched Areas', large_image)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

